#!/bin/bash

# AI Booster 配置更新脚本
# 用于更新代理服务器地址和其他配置

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m'

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查root权限
check_root() {
    if [ "$EUID" -ne 0 ]; then
        print_error "请使用root权限运行此脚本: sudo bash $0"
        exit 1
    fi
}

# 更新代理服务器地址
update_proxy_server() {
    local new_proxy="$1"
    
    if [ -z "$new_proxy" ]; then
        print_error "请提供新的代理服务器地址"
        echo "使用方法: $0 update-proxy <IP:PORT>"
        echo "示例: $0 update-proxy 192.168.1.100:8080"
        exit 1
    fi
    
    print_info "正在更新代理服务器地址为: $new_proxy"
    
    # 更新/etc/ai_booster脚本中的默认代理地址
    if [ -f "/etc/ai_booster" ]; then
        sed -i "s|http://10.205.32.239:9090|http://$new_proxy|g" /etc/ai_booster
        print_success "ai_booster脚本中的代理地址已更新"
    fi
    
    # 更新环境变量
    if grep -q "PUHUI_SMARTML_PROXY_HTTP" /etc/environment; then
        sed -i "s|PUHUI_SMARTML_PROXY_HTTP=.*|PUHUI_SMARTML_PROXY_HTTP=$new_proxy|" /etc/environment
    else
        echo "PUHUI_SMARTML_PROXY_HTTP=$new_proxy" >> /etc/environment
    fi
    print_success "环境变量已更新"
    
    print_success "代理服务器地址更新完成: $new_proxy"
    print_warning "请重新启动终端或重新登录以使更改生效"
}

# 查看当前配置
show_current_config() {
    echo "========================================="
    echo "         当前 AI Booster 配置"
    echo "========================================="
    echo ""
    
    if [ -f "/etc/ai_booster" ]; then
        echo "ai_booster脚本: 已安装"
        
        # 提取当前代理地址
        local current_proxy=$(grep -o 'http://[^"]*' /etc/ai_booster | head -1 | sed 's|http://||')
        echo "默认代理地址: $current_proxy"
    else
        echo "ai_booster脚本: 未安装"
    fi
    
    if grep -q "PUHUI_SMARTML_PROXY_HTTP" /etc/environment; then
        local env_proxy=$(grep "PUHUI_SMARTML_PROXY_HTTP" /etc/environment | cut -d'=' -f2)
        echo "环境变量代理: $env_proxy"
    else
        echo "环境变量代理: 未设置"
    fi
    
    echo ""
    echo "当前代理状态:"
    if [ -n "$http_proxy" ] || [ -n "$https_proxy" ]; then
        echo "  HTTP代理: $http_proxy"
        echo "  HTTPS代理: $https_proxy"
        echo "  状态: 已启用"
    else
        echo "  状态: 未启用"
    fi
    
    echo ""
    echo "========================================="
}

# 修复权限
fix_permissions() {
    print_info "正在修复文件权限..."
    
    if [ -f "/etc/ai_booster" ]; then
        chmod +x /etc/ai_booster
        print_success "ai_booster脚本权限已修复"
    fi
    
    if [ -f "/etc/profile.d/01-locale-fix.sh" ]; then
        chmod +x /etc/profile.d/01-locale-fix.sh
    fi
    
    if [ -f "/etc/profile.d/999-smartml.sh" ]; then
        chmod +x /etc/profile.d/999-smartml.sh
    fi
    
    if [ -f "/etc/profile.d/99-ai-booster-aliases.sh" ]; then
        chmod +x /etc/profile.d/99-ai-booster-aliases.sh
    fi
    
    if [ -f "/usr/local/bin/uninstall_ai_booster.sh" ]; then
        chmod +x /usr/local/bin/uninstall_ai_booster.sh
    fi
    
    print_success "所有文件权限已修复"
}

# 重新安装
reinstall() {
    print_warning "正在重新安装 AI Booster..."
    
    # 检查是否存在安装脚本
    if [ -f "./install_ai_booster.sh" ]; then
        bash ./install_ai_booster.sh
    else
        print_error "找不到安装脚本 install_ai_booster.sh"
        print_info "请确保安装脚本在当前目录下"
        exit 1
    fi
}

# 添加自定义无代理域名
add_no_proxy_domain() {
    local domain="$1"
    
    if [ -z "$domain" ]; then
        print_error "请提供要添加的域名"
        echo "使用方法: $0 add-no-proxy <domain>"
        echo "示例: $0 add-no-proxy example.com"
        exit 1
    fi
    
    print_info "正在添加无代理域名: $domain"
    
    if [ -f "/etc/ai_booster" ]; then
        # 在no_proxy列表中添加新域名
        sed -i "s|no_proxy=localhost|no_proxy=localhost,$domain|" /etc/ai_booster
        print_success "域名 $domain 已添加到无代理列表"
    else
        print_error "ai_booster脚本不存在，请先安装"
        exit 1
    fi
}

# 显示帮助信息
show_help() {
    echo "AI Booster 配置管理脚本"
    echo ""
    echo "使用方法:"
    echo "  $0 <command> [参数]"
    echo ""
    echo "可用命令:"
    echo "  config                 - 查看当前配置"
    echo "  update-proxy <IP:PORT> - 更新代理服务器地址"
    echo "  add-no-proxy <domain>  - 添加无代理域名"
    echo "  fix-permissions        - 修复文件权限"
    echo "  reinstall              - 重新安装"
    echo "  help                   - 显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 config"
    echo "  $0 update-proxy 192.168.1.100:8080"
    echo "  $0 add-no-proxy example.com"
    echo "  $0 fix-permissions"
}

# 主函数
main() {
    case "$1" in
        "config")
            show_current_config
            ;;
        "update-proxy")
            check_root
            update_proxy_server "$2"
            ;;
        "add-no-proxy")
            check_root
            add_no_proxy_domain "$2"
            ;;
        "fix-permissions")
            check_root
            fix_permissions
            ;;
        "reinstall")
            check_root
            reinstall
            ;;
        "help"|"-h"|"--help"|"")
            show_help
            ;;
        *)
            print_error "未知命令: $1"
            echo ""
            show_help
            exit 1
            ;;
    esac
}

# 运行主函数
main "$@"
